import pandas as pd
import plotly.graph_objects as go
## 获取数据


import colorsys
import random

def generate_colors(num_colors, mode='hex', color_space='hsv'):
    """
    生成指定数量的颜色
    
    参数:
        num_colors: 所需颜色数量
        mode: 颜色输出模式，'hex' 或 'rgb'
        color_space: 颜色生成方式，'hsv'（均匀分布）或 'random'（随机生成）
    
    返回:
        list: 颜色字符串列表，格式根据mode参数而定
    """
    if num_colors <= 0:
        return []
    
    colors = []
    
    if color_space == 'hsv':
        # 在HSV颜色空间均匀生成颜色（色调均匀分布，饱和度和明度固定）
        for i in range(num_colors):
            hue = i / num_colors  # 色调从0到1均匀分布
            saturation = 0.7
            value = 0.9
            r, g, b = colorsys.hsv_to_rgb(hue, saturation, value)
            r = int(r * 255)
            g = int(g * 255)
            b = int(b * 255)
            if mode == 'hex':
                # 转换为十六进制格式
                colors.append(f'#{r:02x}{g:02x}{b:02x}')
            else:
                # RGB格式
                colors.append(f'rgb({r},{g},{b})')
    else:
        # 随机生成颜色
        for _ in range(num_colors):
            r = random.randint(0, 255)
            g = random.randint(0, 255)
            b = random.randint(0, 255)
            if mode == 'hex':
                colors.append(f'#{r:02x}{g:02x}{b:02x}')
            else:
                colors.append(f'rgb({r},{g},{b})')
    
    return colors



path = f"./data/全行业用电量季度同比与GDP季度同比.xlsx"


df_sample = pd.read_excel(path, index_col=0).round(2)
df_sample.columns = ['power', 'gdp']



import pandas as pd
import plotly.graph_objects as go
import plotly.express as px  # 用于获取颜色序列


def plot_gdp_power_scatter(df, title="GDP与用电量关联分析"):
    """
    绘制GDP与用电量的散点图，并添加表示时间序列方向的箭头，不同年份的散点用不同颜色区分
    
    参数:
    df: DataFrame, 索引为日期格式，包含'power'和'gdp'两列
    title: 图表标题
    
    返回:
    fig: plotly图形对象
    """
    
    # 确保索引是日期格式
    df = df.copy()
    df.index = pd.to_datetime(df.index)
    
    # 按日期排序
    df = df.sort_index()
    
    # 创建图表
    fig = go.Figure()
    
    year_data = df
    # 先添加所有散点（按年份区分颜色）
    fig.add_trace(go.Scatter(
        x=year_data['power'],
        y=year_data['gdp'],
        mode='markers',
        # name=str(year),  # 用年份作为图例名称
        marker=dict(
            size=8,
            color='blue',
            opacity=0.7
        ),
        text=[f"日期: {idx.strftime('%Y-%m')}<br>用电量: {x:.2f}<br>GDP: {y:.2f}" 
                for idx, x, y in zip(year_data.index, year_data['power'], year_data['gdp'])],
        hoverinfo='text'
    ))
    
    # 设置布局
    fig.update_layout(
        title=title,
        xaxis_title="用电量",
        yaxis_title="GDP",
        hovermode="closest",
        height=600,
        xaxis=dict(
            zeroline=True,
            zerolinecolor='white',
            zerolinewidth=1
        ),
        yaxis=dict(
            zeroline=True,
            zerolinecolor='white',
            zerolinewidth=1
        )
    )
    return fig


def plot_gdp_power_scatter_with_arrows(df, title="GDP与用电量关联分析"):
    """
    绘制GDP与用电量的散点图，并添加表示时间序列方向的箭头，不同年份的散点用不同颜色区分
    
    参数:
    df: DataFrame, 索引为日期格式，包含'power'和'gdp'两列
    title: 图表标题
    
    返回:
    fig: plotly图形对象
    """
    
    # 确保索引是日期格式
    df = df.copy()
    df.index = pd.to_datetime(df.index)
    
    # 按日期排序
    df = df.sort_index()
    
    # 获取所有唯一年份并分配颜色
    years = df.index.year.unique()
    # 使用plotly的内置颜色序列，数量不够时会循环使用
    hex_colors = generate_colors(len(years), mode='hex', color_space='hsv')
    # colors = px.colors.qualitative.Plotly[:len(years)]
    year_color_map = {year: color for year, color in zip(years, hex_colors)}
    
    # 创建图表
    fig = go.Figure()
    
    # 先添加所有散点（按年份区分颜色）
    for year in years:
        year_data = df[df.index.year == year]
        fig.add_trace(go.Scatter(
            x=year_data['power'],
            y=year_data['gdp'],
            mode='markers',
            name=str(year),  # 用年份作为图例名称
            marker=dict(
                size=8,
                color=year_color_map[year],
                opacity=0.7
            ),
            text=[f"日期: {idx.strftime('%Y-%m')}<br>用电量: {x:.2f}<br>GDP: {y:.2f}" 
                  for idx, x, y in zip(year_data.index, year_data['power'], year_data['gdp'])],
            hoverinfo='text'
        ))
    
    # 再添加连线（保持原有颜色逻辑）
    for i in range(len(df) - 1):
        ddf = df.iloc[i:i+2]
        from_x = ddf['power'].iloc[0]
        from_y = ddf['gdp'].iloc[0]
        to_x = ddf['power'].iloc[1]
        to_y = ddf['gdp'].iloc[1]

        # 保持原有连线颜色逻辑（正相关为绿色，负相关为红色）
        if (to_x - from_x) * (to_y - from_y) > 0:
            line_color = 'green'
        else:
            line_color = 'red'

        fig.add_trace(go.Scatter(
            x=ddf['power'],
            y=ddf['gdp'],
            mode='lines',
            showlegend=False,  # 连线不显示在图例中
            line=dict(
                color=line_color,
                width=1,
                dash='dot'
            ),
            hoverinfo='none'  # 连线不显示hover信息
        ))
    

    # 设置布局
    fig.update_layout(
        title=title,
        xaxis_title="用电量",
        yaxis_title="GDP",
        hovermode="closest",
        height=600,
        xaxis=dict(
            zeroline=True,
            zerolinecolor='white',
            zerolinewidth=1
        ),
        yaxis=dict(
            zeroline=True,
            zerolinecolor='white',
            zerolinewidth=1
        )
    )
    
    return fig

# plot_gdp_power_scatter_with_arrows(df_sample)
plot_gdp_power_scatter(df_sample)